package com.rain.rainspringboot3.encryption.strategy;

import com.rain.rainspringboot3.encryption.exception.DecryptionException;
import com.rain.rainspringboot3.encryption.exception.EncryptionException;
import jakarta.annotation.PostConstruct;
import org.apache.commons.lang3.ObjectUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.crypto.Cipher;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

import static com.rain.rainspringboot3.encryption.service.EncryptionServiceImpl.KEY_CACHE_MAP;

/**
 * RSA加密策略实现
 * 使用RSA算法提供加密和解密功能
 */
@Component
public class RSAEncryptionStrategy implements EncryptionStrategy {

    private static final String ALGORITHM = "RSA/ECB/PKCS1Padding";
    private static final int KEY_SIZE = 2048;

    @Value("${encryption.strategies.rsa.public-key}")
    private String publicKeyStr;

    @Value("${encryption.strategies.rsa.private-key}")
    private String privateKeyStr;

    private PublicKey publicKey;
    private PrivateKey privateKey;

    @Override
    public String encrypt(String content, String key) {
        try {
            // 如果传入了key，使用传入的key，否则使用配置的key
            PublicKey usePublicKey = publicKey;
            if (StringUtils.isNotBlank(key) && ObjectUtils.isEmpty(usePublicKey)) {
                String[] keys = key.split("\\|");
                usePublicKey = getPublicKeyFromString(keys[0]);
            }

            if (usePublicKey == null) {
                throw new EncryptionException("RSA public key is not available");
            }

            Cipher cipher = Cipher.getInstance(ALGORITHM);
            cipher.init(Cipher.ENCRYPT_MODE, usePublicKey);

            // RSA加密有长度限制，需要分段加密
            byte[] data = content.getBytes(StandardCharsets.UTF_8);
            int inputLen = data.length;
            int maxEncryptBlock = 245; // 2048位密钥，PKCS1Padding模式下最大加密块为245字节

            try (java.io.ByteArrayOutputStream out = new java.io.ByteArrayOutputStream()) {
                int offSet = 0;
                byte[] cache;
                int i = 0;

                // 对数据分段加密
                while (inputLen - offSet > 0) {
                    if (inputLen - offSet > maxEncryptBlock) {
                        cache = cipher.doFinal(data, offSet, maxEncryptBlock);
                    } else {
                        cache = cipher.doFinal(data, offSet, inputLen - offSet);
                    }
                    out.write(cache, 0, cache.length);
                    i++;
                    offSet = i * maxEncryptBlock;
                }

                byte[] encryptedData = out.toByteArray();
                return Base64.getEncoder().encodeToString(encryptedData);
            }
        } catch (Exception e) {
            throw new EncryptionException("RSA encryption failed", e);
        }
    }

    @Override
    public String decrypt(String content, String key) {
        try {
            // 如果传入了key，使用传入的key，否则使用配置的key
            PrivateKey usePrivateKey = privateKey;
            if (StringUtils.isNotBlank(key) && ObjectUtils.isEmpty(usePrivateKey)) {
                String[] keys = key.split("\\|");
                usePrivateKey = getPrivateKeyFromString(keys[1]);
            }

            if (usePrivateKey == null) {
                throw new DecryptionException("RSA private key is not available");
            }

            Cipher cipher = Cipher.getInstance(ALGORITHM);
            cipher.init(Cipher.DECRYPT_MODE, usePrivateKey);

            // RSA解密有长度限制，需要分段解密
            byte[] data = Base64.getDecoder().decode(content);
            int inputLen = data.length;
            int maxDecryptBlock = 256; // 2048位密钥，PKCS1Padding模式下最大解密块为256字节

            try (java.io.ByteArrayOutputStream out = new java.io.ByteArrayOutputStream()) {
                int offSet = 0;
                byte[] cache;
                int i = 0;

                // 对数据分段解密
                while (inputLen - offSet > 0) {
                    if (inputLen - offSet > maxDecryptBlock) {
                        cache = cipher.doFinal(data, offSet, maxDecryptBlock);
                    } else {
                        cache = cipher.doFinal(data, offSet, inputLen - offSet);
                    }
                    out.write(cache, 0, cache.length);
                    i++;
                    offSet = i * maxDecryptBlock;
                }

                byte[] decryptedData = out.toByteArray();
                return new String(decryptedData, StandardCharsets.UTF_8);
            }
        } catch (Exception e) {
            throw new DecryptionException("RSA decryption failed", e);
        }
    }

    @Override
    public String generateKey() {
        try {
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
            keyPairGenerator.initialize(KEY_SIZE);
            KeyPair keyPair = keyPairGenerator.generateKeyPair();

            // 存储生成的密钥
            publicKey = keyPair.getPublic();
            privateKey = keyPair.getPrivate();

            // 返回公钥和私钥的Base64编码字符串，用"|"分隔
            String publicKeyStr = Base64.getEncoder().encodeToString(publicKey.getEncoded());
            String privateKeyStr = Base64.getEncoder().encodeToString(privateKey.getEncoded());

            return publicKeyStr + "|" + privateKeyStr;
        } catch (Exception e) {
            throw new EncryptionException("Failed to generate RSA key pair", e);
        }
    }


    /**
     * 从字符串加载公钥
     */
    private PublicKey getPublicKeyFromString(String keyStr) {
        try {
            // 提取Base64内容部分
            String cleanedKey = extractBase64Content(keyStr);

            byte[] keyBytes = Base64.getDecoder().decode(cleanedKey);
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            return keyFactory.generatePublic(keySpec);
        } catch (Exception e) {
            e.printStackTrace();
            throw new EncryptionException("Failed to load public key from string", e);
        }
    }

    /**
     * 提取Base64内容，移除PEM头尾和非法字符
     */
    private String extractBase64Content(String keyStr) {
        if (keyStr == null || keyStr.isEmpty()) {
            throw new IllegalArgumentException("Key string is null or empty");
        }

        // 移除PEM头尾标识
        String content = keyStr
                .replaceAll("-----BEGIN (.*)-----", "")
                .replaceAll("-----END (.*)-----", "")
                .replaceAll("\\r\\n", "")
                .replaceAll("\\n", "")
                .replaceAll("\\s", "");

        // 处理URL安全的Base64字符
        return content.replace("-", "+").replace("_", "/");
    }



    /**
     * 从字符串加载私钥
     */
    /**
     * 从字符串加载私钥
     */
    private PrivateKey getPrivateKeyFromString(String keyStr) {
        try {
            // 处理PEM格式的私钥字符串
            String privateKeyContent = keyStr
                    .replace("-----BEGIN PRIVATE KEY-----", "")
                    .replace("-----END PRIVATE KEY-----", "")
                    .replace("-----BEGIN RSA PRIVATE KEY-----", "")
                    .replace("-----END RSA PRIVATE KEY-----", "")
                    .replaceAll("\\s", "")
                    .replace("-", "+")
                    .replace("_", "/");

            byte[] keyBytes = Base64.getDecoder().decode(privateKeyContent);
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            return keyFactory.generatePrivate(keySpec);
        } catch (Exception e) {
            e.printStackTrace();
            throw new EncryptionException("Failed to load private key from string", e);
        }
    }


    /**
     * 初始化方法，从配置加载密钥
     */
    @PostConstruct
    public void init() {
        if (publicKeyStr != null && !publicKeyStr.trim().isEmpty()) {
            publicKey = getPublicKeyFromString(publicKeyStr);
        }

        if (privateKeyStr != null && !privateKeyStr.trim().isEmpty()) {
            privateKey = getPrivateKeyFromString(privateKeyStr);
        }
    }

    @PostConstruct
    public void saveKey(){
        KEY_CACHE_MAP.put("RSA",publicKeyStr);
    }
}